import cv2
import numpy as np
import glob

# 1. 设置棋盘格参数
pattern_size = (11, 8)  # 内部角点个数
square_size = 0.025  # 每格大小（单位：米）

# 2. 准备真实世界坐标点（z=0平面）
objp = np.zeros((pattern_size[0]*pattern_size[1], 3), np.float32)
objp[:, :2] = np.indices(pattern_size).T.reshape(-1, 2)
objp *= square_size

objpoints = []  # 3D 点
imgpoints = []  # 2D 点

# 3. 读取标定图像
# images = glob.glob('calib_images/*.jpg')
images = glob.glob('images_tmp/*.jpg')
for fname in images:
    img = cv2.imread(fname)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    ret, corners = cv2.findChessboardCorners(gray, pattern_size, None)
    if ret:
        objpoints.append(objp)
        corners2 = cv2.cornerSubPix(gray, corners, (11,11), (-1,-1),
                                     criteria=(cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001))
        imgpoints.append(corners2)
        cv2.drawChessboardCorners(img, pattern_size, corners2, ret)
        cv2.imshow('calib', img)
        cv2.waitKey(100)

cv2.destroyAllWindows()

# 4. 相机内参标定

if len(objpoints) > 0 and len(imgpoints) > 0:
    ret, mtx, dist, rvecs, tvecs = cv2.calibrateCamera(objpoints, imgpoints, gray.shape[::-1], None, None)
else:
    print("No valid image points were found. Make sure the calibration pattern is detected in the input images.")

# 5. 单应矩阵标定（选取一张图像与其角点）
img = cv2.imread(images[0])
gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
ret, corners = cv2.findChessboardCorners(gray, pattern_size, None)
if ret:
    img_pts = corners.reshape(-1, 2)
    obj_pts = objp[:, :2]  # 取xy平面
    H, _ = cv2.findHomography(img_pts, obj_pts)

    # 打印相机内参和畸变系数
    print("相机内参矩阵:\n", mtx)
    print("畸变系数:\n", dist)
    print("单应矩阵:\n", H)
    # 6. 计算相机外参（旋转矩阵和平移向量）
    rvec, tvec = cv2.Rodrigues(rvecs[0])
    print("旋转向量:\n", rvec)
    print("平移向量:\n", tvec)
    # 7. 显示标定结果
    h, w = img.shape[:2]
    newcameramtx, roi = cv2.getOptimalNewCameraMatrix(mtx, dist, (w, h), 1, (w, h))
    dst = cv2.undistort(img, mtx, dist, None, newcameramtx)
    x, y, w, h = roi
    dst = dst[y:y+h, x:x+w]
    cv2.imshow('undistorted', dst)
    cv2.waitKey(0)
    cv2.destroyAllWindows()

    # 8. 计算相机的FOV（视场角）
    fov_x = 2 * np.arctan(w / (2 * newcameramtx[0, 0])) * 180 / np.pi
    fov_y = 2 * np.arctan(h / (2 * newcameramtx[1, 1])) * 180 / np.pi
    print("水平视场角: {:.2f}°".format(fov_x))
    print("垂直视场角: {:.2f}°".format(fov_y))
    # 9. 计算相机的焦距
    focal_length = np.mean([mtx[0, 0], mtx[1, 1]])
    print("焦距: {:.2f} 像素".format(focal_length))
    # 10. 计算相机的畸变系数
    print("畸变系数:\n", dist)
    # 11. 计算相机的主点坐标
    principal_point = (mtx[0, 2], mtx[1, 2])
    print("主点坐标: (x: {:.2f}, y: {:.2f})".format(principal_point[0], principal_point[1]))
    # 12. 计算相机的视场范围
    fov_x_rad = np.radians(fov_x)
    fov_y_rad = np.radians(fov_y)
    view_width = 2 * focal_length * np.tan(fov_x_rad / 2)


    # 保存标定参数
    np.savez('calibration_data.npz', mtx=mtx, dist=dist, homography=H)
    print("✅ 标定完成，已保存 calibration_data.npz")
else:
    print("❌ 未找到角点，无法生成 homography")
